from langchain.memory import ConversationBufferWindowMemory
from langchain_openai import AzureChatOpenAI
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.llms import OpenAI
from docx import Document
# from dotenv import load_dotenv

import os


# load_dotenv()

os.environ["AZURE_OPENAI_ENDPOINT"] = "https://chatgpt-test-0001.openai.azure.com/"
os.environ["AZURE_OPENAI_API_KEY"] = "90643b9c4bee421bba17d72b748ad446"

# 创建一个函数生成毕业设计文档
def generate_project_document(titles, filename="毕业设计.docx"):
    document = Document()

    # 添加文档标题
    document.add_heading('基于云计算的智能校园网设计与实现', 0)

    # 遍历章节标题
    for index, title in enumerate(titles):
        document.add_heading(f'第{index + 1}章 {title}', level=1)

        # 利用OpenAI生成章节内容
        chapter_content = generate_content(title)

        # 将生成的内容添加到文档
        document.add_paragraph(chapter_content.content)

    # 保存文档
    document.save(filename)
    print(f"文档已生成并保存为: {filename}")

store = {}  # memory is maintained outside the chain

def get_session_history(session_id: str) -> InMemoryChatMessageHistory:
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
        return store[session_id]

    memory = ConversationBufferWindowMemory(
        chat_memory=store[session_id],
        k=3,
        return_messages=True,
    )
    assert len(memory.memory_variables) == 1
    key = memory.memory_variables[0]
    messages = memory.load_memory_variables({})[key]
    store[session_id] = InMemoryChatMessageHistory(messages=messages)
    return store[session_id]

# 使用LangChain和OpenAI生成文本
def generate_content(title):
    # 在此处使用langchain的ConversationChain进行对话生成

    prompt = ChatPromptTemplate.from_messages([
        ("system", "You're an assistant who's good at writing"),
        MessagesPlaceholder(variable_name="history"),
        ("human", "请围绕以下主题生成一段详细的内容：{title}，并且要求字数不少于500字。同时需求分析、设计方案和功能实现部分需要有配图进行介绍和概括，"),
    ])

    chain = prompt | AzureChatOpenAI(
        azure_deployment="testMini",  # or your deployment
        api_version="2024-06-01",  # or your api version
        temperature=0,
        max_tokens=None,
        timeout=None,
        max_retries=2,)

    chain_with_history = RunnableWithMessageHistory(
        chain,
        # Uses the get_by_session_id function defined in the example
        # above.
        get_session_history,
        input_messages_key="title",
        history_messages_key="history",
    )

    # 利用OpenAI生成文本
    return chain_with_history.invoke(
        {"title": title},
        config={"configurable": {"session_id": "foo"}})


if __name__ == "__main__":
    # 定义文档的章节标题
    titles = [
        "引言"
    ]

    # 生成Word文档
    generate_project_document(titles)
